#include "pch.h"
#include <memory>
#include "Service.h"
#include "ImportData.h"
#include "Util.h"
#include "SyncData.h"
#include "../Config/Vault.h"

Service::Service(const std::string &config_file, std::shared_ptr<EnclaveProvider> enclave_provider)
    : config_(config_file), enclave_provider_(enclave_provider)
{
	LOG(INFO) << "Hana host: " << config_.database_.host_ << ", port: " << config_.database_.port_ 
		<< ", database: " << config_.database_.database_ << ", schema: " << config_.database_.schema_;

	enclave_providers_map_[enclave_provider->id_] = enclave_provider_;
}


Service::~Service()
{
}

void Service::start()
{
    LOG(INFO) << "Start grpc service";

    char server_address_buf[64] = { 0 };
    sprintf_s(server_address_buf, 64, "0.0.0.0:%d", config_.port_);

    grpc::ServerBuilder builder;
    builder.AddListeningPort(server_address_buf, grpc::InsecureServerCredentials());
    builder.RegisterService(this);
	builder.SetMaxReceiveMessageSize(INT_MAX);
	builder.SetMaxSendMessageSize(INT_MAX);

    std::unique_ptr<grpc::Server> server(builder.BuildAndStart());

    LOG(INFO) << "Server listening on " << server_address_buf;

    server->Wait();
}

::grpc::Status Service::ImportData(::grpc::ServerContext* context, const ::enclave::ImportDataRequest* request, ::grpc::ServerWriter< ::enclave::ImportDataResponse>* writer)
{
	LOG(INFO) << "Import data rqueset, " << "host: " << request->host() << ", port: " << request->port();

    ::ImportData import_data = ::ImportData(enclave_provider_, config_);

    import_data.importData(request, writer);

    return ::grpc::Status::OK;
}

::grpc::Status Service::Search(::grpc::ServerContext * context, const::enclave::searchRequest * request, ::grpc::ServerWriter<::enclave::Rows>* writer)
{
	LOG(INFO) << "Search request: " << request->sql();
	
    std::vector<std::string> args;
    for (int i = 0; i < request->values_size(); i++)
    {
		auto value = request->values(i).value();
		
		LOG(INFO) << value;

        if (request->values(i).type() == enclave::searchRequest::valueType::searchRequest_valueType_intType)
        {
            args.push_back(value);
        }
        else
        {
			if (value.empty())
			{
				args.push_back(value);
			}
			else
			{
				int buf_len = (int)value.length() * 4;
				std::unique_ptr<char> buf = std::unique_ptr<char>(new char[buf_len]);
				memset(buf.get(), 0, buf_len);

				int len = enclave_provider_->encrypt(value.c_str(), (int)value.length(), buf.get(), buf_len);

				if (len == 0)
				{
					LOG(INFO) << "Encrypt agts faild";
					return ::grpc::Status::CANCELLED;
				}

				args.push_back(std::string(buf.get(), len).c_str());
			}
        }
    }

    HanaData hanaData = HanaData(config_.database_);

    IData::Fields fields;
	int num_rows = 0;
    std::vector<Field> columns;

	LOG(INFO) << "Selecting ...";

    hanaData.select(request->sql(), args, fields, num_rows, &columns);

	LOG(INFO) << "Result number rows = " << num_rows;

    enclave::Rows response;

	char temp[256] = { 0 };

    for (size_t i = 0; i < columns.size(); ++i)
    {
        if (std::find(ImportData::columns_ex_.begin(), ImportData::columns_ex_.end(), columns[i]) != ImportData::columns_ex_.end()
			|| columns[i].name_ == "SEQ")
        {
            continue;
        }

        auto column = response.add_columns();

//		memset(temp, 0, sizeof(temp));
//		enclave_provider_->decrypt(columns[i].name_.c_str(), (int)columns[i].name_.length(), temp, (int)sizeof(temp));
//
//		columns[i].name_ = temp;

        column->set_field(columns[i].name_);
    }

    response.set_total(num_rows);

    int retval = 0;
    int buf_len = 5120;
    char *buf = new char[buf_len];
	int num_fields = (int)columns.size();
	fields.resetPos();
    for (size_t i = 0; i < num_rows; ++i)
    {
        auto row = response.add_rows();
        auto values = row->mutable_values();

        for (size_t j = 0; j < num_fields; ++j)
        {
			int field_len = fields.field_len(int(i * num_fields + j));

            if (std::find(ImportData::columns_ex_.begin(), ImportData::columns_ex_.end(), columns[j]) != ImportData::columns_ex_.end())
            {
				fields.buf_pos_ += field_len;
                continue;
            }
            
			if (columns[j].type_ == "VARCHAR")
			{
				memset(buf, 0, buf_len);
				enclave_provider_->decrypt(fields.buf() + fields.buf_pos_, field_len, buf, buf_len);
				*values += buf;
				row->add_value_lens((int)strlen(buf));
			}
			else 
			{
				*values += std::string(fields.buf() + fields.buf_pos_, field_len);
				row->add_value_lens(field_len);
			}

			fields.buf_pos_ += field_len;
        }

        if (response.rows_size() == 200)
        {
            response.set_pos(int(i + 1 - response.rows_size()));
            response.set_length(response.rows_size());
            writer->Write(response);

            response.clear_rows();
        }
    }
    delete[]buf;

    response.set_pos(int(num_rows - response.rows_size()));
    response.set_length(response.rows_size());

    grpc::WriteOptions write_options;
    writer->WriteLast(response, write_options);

    return ::grpc::Status::OK;
}

::grpc::Status Service::SyncCmd(::grpc::ServerContext* context, const ::enclave::SyncCmdRequest* request, ::enclave::SyncCmdResponse* response)
{
    LOG(INFO) << "SyncCmd , valutId: " << request->vaultid() << "host: " << request->host() << "port: " << request->port();

    char str_guid[64] = { 0 };
    GUID guid;

    if (CoCreateGuid(&guid))
    {
        LOG(ERROR) << "Create guid error";
        return ::grpc::Status::CANCELLED;
    }
    _snprintf_s(str_guid, sizeof(str_guid),
        "%08X-%04X-%04x-%02X%02X-%02X%02X%02X%02X%02X%02X",
        guid.Data1, guid.Data2, guid.Data3,
        guid.Data4[0], guid.Data4[1], guid.Data4[2],
        guid.Data4[3], guid.Data4[4], guid.Data4[5],
        guid.Data4[6], guid.Data4[7]);

	char address[256] = { 0 };
	sprintf_s(address, sizeof(address), "%s:%d", request->host().c_str(), request->port());
	auto channel = grpc::CreateChannel(address, grpc::InsecureChannelCredentials());
	if (channel.get() == nullptr)
	{
		LOG(ERROR) << "Create channel error.";
		return ::grpc::Status::CANCELLED;
	}

	auto stub = enclave::EnclaveServer::NewStub(channel);

	LOG(INFO) << "Sync table info, sync id = " << str_guid;

	::SyncData sync_data(enclave_provider_, &config_);

	{
		HanaData hana = HanaData(config_.database_);

		if (hana.hdbc_ == nullptr)
		{
			LOG(ERROR) << "Hana database connect error";
			return ::grpc::Status::CANCELLED;
		}

		enclave::SyncNotifyRequest sync_notify_request;
		sync_notify_request.set_syncid(str_guid);
		sync_notify_request.set_vaultid(enclave_provider_->id_);

		std::vector<std::string> tables;
		hana.showTables(tables);

		for (auto iter = tables.begin(); iter != tables.end(); iter++)
		{
			std::vector<Field> columns;
			hana.showColumns(*iter, columns);

			auto table = sync_notify_request.add_tables();

			table->set_name(*iter);
			for (auto column_iter = columns.begin(); column_iter != columns.end(); column_iter++)
			{
				auto column = table->add_columns();

				Util::IColumnToGColumn(*column_iter, *column, true, true);
			}

			sync_data.tables_columns_[*iter] = columns;

			LOG(INFO) << *iter << " columns count = " << columns.size();
		}

		LOG(INFO) << "Send syncNotify";

		grpc::ClientContext sync_notify_context;
		enclave::SyncTablesDatas sync_notify_response;
		auto sync_notify_status = stub->SyncNotify(&sync_notify_context, sync_notify_request, &sync_notify_response);

		if (sync_notify_status.ok() == false)
		{
			LOG(ERROR) << sync_notify_status.error_message();
			return grpc::Status::CANCELLED;
		}

		enclave::SyncTablesDatas sync_info_request;
		sync_info_request.set_syncid(str_guid);
		sync_info_request.set_vaultid(enclave_provider_->id_);

		LOG(INFO) << "Generator sync task ...";
		sync_data.generatorTask(hana, tables, sync_notify_response, sync_info_request);
	
		LOG(INFO) << "Send syncInfo";
		grpc::ClientContext sync_info_context;
		enclave::SyncInfoResponse sync_info_response;
		auto sync_info_status = stub->SyncInfo(&sync_info_context, sync_info_request, &sync_info_response);

		if (sync_info_status.ok() == false)
		{
			LOG(ERROR) << sync_info_status.error_message();
			return grpc::Status::CANCELLED;
		}
	}

	Vault vault;
	vault.id_ = request->vaultid();
	vault.host_ = request->host();
	vault.port_ = request->port();

	LOG(INFO) << "Start sync ...";
    auto enclave_provider = get_enclave_porvider(vault.id_);
	sync_data.syncData(str_guid, vault, enclave_provider);

	LOG(INFO) << "Send sync finished";
	grpc::ClientContext sync_context;
	enclave::SyncFinishedRequest sync_finished_request;
	enclave::SyncFinishedResponse sync_finished_response;
	sync_finished_request.set_syncid(str_guid);
	sync_finished_request.set_vaultid(enclave_provider_->id_);
	LOG(INFO) << "data send done, waiting for finishe";
	auto sync_finished_finished = stub->SyncFinished(&sync_context, sync_finished_request, &sync_finished_response);
	LOG(INFO) << "sync_finished_finished = " << sync_finished_finished.error_code() << ", messge = " << sync_finished_finished.error_message();

    LOG(INFO) << "Sync complete";
    response->set_vaultid(enclave_provider_->id_);

    return ::grpc::Status::OK;
}

::grpc::Status Service::SyncNotify(::grpc::ServerContext* context, const ::enclave::SyncNotifyRequest* request, ::enclave::SyncTablesDatas* response)
{
    LOG(INFO) << "SyncNotify";

    std::shared_ptr<::SyncData> sync_data = std::shared_ptr<::SyncData>(new ::SyncData(enclave_provider_, &config_));

    ::SyncData::sync_datas_map[request->syncid()] = sync_data;

    sync_data->onNotifySync(request, response);

    response->set_vaultid(enclave_provider_->id_);
    response->set_syncid(request->syncid());

    return ::grpc::Status::OK;
}

::grpc::Status Service::SyncInfo(::grpc::ServerContext* context, const ::enclave::SyncTablesDatas* request, ::enclave::SyncInfoResponse* response)
{
    LOG(INFO) << "SyncInfo";

    if (::SyncData::sync_datas_map.find(request->syncid()) == ::SyncData::sync_datas_map.end())
    {
        return ::grpc::Status(::grpc::StatusCode::CANCELLED, "Can not find sync id" + request->syncid());
    }

    std::shared_ptr<::SyncData> sync_data = ::SyncData::sync_datas_map[request->syncid()];

    sync_data->onSyncInfo(request);

    response->set_vaultid(enclave_provider_->id_);

    return ::grpc::Status::OK;
}

::grpc::Status Service::SyncData(::grpc::ServerContext* context, ::grpc::ServerReader< ::enclave::SyncTableRows>* reader, ::enclave::SyncDataResponse* response)
{
    std::shared_ptr<enclave::SyncTableRows> sync_table_rows = std::shared_ptr<enclave::SyncTableRows>(new enclave::SyncTableRows);

    while (reader->Read(sync_table_rows.get()))
    {
        if (::SyncData::sync_datas_map.find(sync_table_rows->syncid()) == ::SyncData::sync_datas_map.end())
        {
			LOG(INFO) << "Can not find sync id" << sync_table_rows->syncid();
            return ::grpc::Status(::grpc::StatusCode::CANCELLED, "Can not find sync id" + sync_table_rows->syncid());
        }

        std::shared_ptr<::SyncData> sync_data = SyncData::sync_datas_map[sync_table_rows->syncid()];

		LOG(INFO) << "Sync Data, start = " << sync_table_rows->rows().pos() << ", end = " << sync_table_rows->rows().pos() + sync_table_rows->rows().length() - 1;;

		sync_data->addSyncDatas(sync_table_rows);

		sync_table_rows = std::shared_ptr<enclave::SyncTableRows>(new enclave::SyncTableRows);
    };

    return ::grpc::Status::OK;
}

::grpc::Status Service::SyncFinished(::grpc::ServerContext* context, const ::enclave::SyncFinishedRequest* request, ::enclave::SyncFinishedResponse* response)
{
	LOG(INFO) << "SyncFinishing";

	if (::SyncData::sync_datas_map.find(request->syncid()) == ::SyncData::sync_datas_map.end())
	{
		return ::grpc::Status(::grpc::StatusCode::CANCELLED, "Can not find sync id" + request->syncid());
	}

	std::shared_ptr<::SyncData> sync_data = ::SyncData::sync_datas_map[request->syncid()];

	sync_data->setSyncFinished(true);

	LOG(INFO) << "SyncFinished";

	return ::grpc::Status::OK;
}

::grpc::Status Service::PubEncrypt(::grpc::ServerContext* context, const ::enclave::Rows* request, ::enclave::Rows* response)
{
	LOG(INFO) << "PubEncrypt";

	int buf_len = 10240;
	std::unique_ptr<char> buf = std::unique_ptr<char>(new char[buf_len]);

	int field_buf_len = 5000;
	std::unique_ptr<char> field_buf = std::unique_ptr<char>(new char[field_buf_len]);

	IData::Fields fields;

	std::vector<Field> columns;

	response->set_tablename(request->tablename());

	response->set_total(request->total());

	response->set_pos(request->pos());

	response->set_length(request->length());

	for (int i = 0; i < request->columns_size(); i++)
	{
		auto column_info = request->columns(i);
		Field column;

		Util::GColumnToIColumn(column_info, column, true, true);

		columns.push_back(column);

		auto response_column = response->add_columns();

		Util::IColumnToGColumn(column, *response_column, true, true);
	}

	for (int i = 0; i < request->rows_size(); i++)
	{
		auto values = request->rows(i).values();

		auto row = response->add_rows();

		int pos = 0;

		for (int j = 0; j < request->rows(i).value_lens_size(); j++)
		{
			int value_len = request->rows(i).value_lens(j);

			if (std::find(ImportData::columns_ex_.begin(), ImportData::columns_ex_.end(), columns[j]) != ImportData::columns_ex_.end())
			{
				fields.add_field(values.c_str() + pos, value_len);

				row->add_value_lens(value_len);
			} 
			else if (columns[j].type_ == "DECIMAL" || columns[j].type_ == "INTEGER")
			{
				fields.add_field(values.c_str() + pos, value_len);

				row->add_value_lens(value_len);
			}
			else 
			{
				if (value_len == 0)
				{
					fields.add_field(nullptr, 0);
					row->add_value_lens(0);
				}
				else
				{
					memset(field_buf.get(), 0, field_buf_len);

					int decrypt_value_len = enclave_provider_->decrypt(values.c_str() + pos, value_len, field_buf.get(), field_buf_len);

					if (decrypt_value_len == 0)
					{
						response->clear_rows();
						LOG(ERROR) << "Decrypt error, values: " << std::string(values.c_str() + pos, value_len);
						return ::grpc::Status(::grpc::StatusCode::CANCELLED, "Decrypt error");
					}

					fields.add_field(field_buf.get(), decrypt_value_len);
					row->add_value_lens(decrypt_value_len);
				}
			}
			
			pos += value_len;
		}
	
		memset(buf.get(), 0, buf_len);

		auto enclave_provider = get_enclave_porvider(request->vaultid());

		int encrypt_len = enclave_provider->public_key_encrypt(fields.buf(), fields.buf_pos_, buf.get(), buf_len);

		if (encrypt_len == 0)
		{
			response->clear_rows();
			LOG(ERROR) << "PubEncrypt error, values: " << std::string(fields.buf(), fields.buf_pos_);
			return ::grpc::Status(::grpc::StatusCode::CANCELLED, "PubEncrypt error");
		}

		fields.clear();

		row->set_values(buf.get(), encrypt_len);
	}

	LOG(INFO) << "PubEncrypt end";
	
	return ::grpc::Status::OK;
}

::grpc::Status Service::PriDecrypt(::grpc::ServerContext* context, const ::enclave::Rows* request, ::enclave::Rows* response)
{
	LOG(INFO) << "PriDecrypt";

	int buf_len = 10240;
	std::unique_ptr<char> buf = std::unique_ptr<char>(new char[buf_len]);

	int field_buf_len = 5000;
	std::unique_ptr<char> field_buf = std::unique_ptr<char>(new char[field_buf_len]);
	
	IData::Fields fields;

	std::vector<Field> columns;

	response->set_tablename(request->tablename());

	response->set_total(request->total());

	response->set_pos(request->pos());

	response->set_length(request->length());

	for (int i = 0; i < request->columns_size(); i++)
	{
		auto column_info = request->columns(i);
		Field column;

		Util::GColumnToIColumn(column_info, column, true, true);

		columns.push_back(column);

		auto response_column = response->add_columns();

		Util::IColumnToGColumn(column, *response_column, true, true);
	}

	for (int i = 0; i < request->rows_size(); i++)
	{
		auto values = request->rows(i).values();

		memset(buf.get(), 0, buf_len);
		int decrypt_len = enclave_provider_->private_key_decrypt(values.c_str(), (int)values.length(), buf.get(), buf_len);

		if (decrypt_len == 0)
		{
			LOG(INFO) << "Private key decrypt error, values = " << values;
			response->clear_rows();
			return ::grpc::Status(::grpc::StatusCode::CANCELLED, "Private key decrypt error");
		}

		auto row = response->add_rows();

		int pos = 0;
		for (int j = 0; j < request->rows(i).value_lens_size(); j++)
		{
			int value_len = request->rows(i).value_lens(j);

			if (std::find(ImportData::columns_ex_.begin(), ImportData::columns_ex_.end(), columns[j]) != ImportData::columns_ex_.end())
			{
				fields.add_field(buf.get() + pos, value_len);
				row->add_value_lens(value_len);
			}
			else if (columns[j].type_ == "DECIMAL" || columns[j].type_ == "INTEGER")
			{
				fields.add_field(buf.get() + pos, value_len);
				row->add_value_lens(value_len);
			}
			else
			{
				if (value_len == 0)
				{
					fields.add_field(nullptr, 0);
					row->add_value_lens(0);
				}
				else
				{
					memset(field_buf.get(), 0, field_buf_len);

					int encrypt_len = enclave_provider_->encrypt(buf.get() + pos, value_len, field_buf.get(), field_buf_len);
					if (encrypt_len == 0)
					{
						LOG(INFO) << "Encrypt error, value = " << std::string(buf.get() + pos);
						response->clear_rows();
						return ::grpc::Status(::grpc::StatusCode::CANCELLED, "Encrypt error");
					}

					fields.add_field(field_buf.get(), encrypt_len);
					row->add_value_lens(encrypt_len);
				}
			}

			pos += value_len;
		}

		row->set_values(fields.buf(), fields.buf_pos_);

		fields.clear();
	}

	LOG(INFO) << "PriDecrypt end";

	return ::grpc::Status::OK;
}

::grpc::Status Service::Encrypt(::grpc::ServerContext* context, const ::enclave::Fields* request, ::enclave::Fields* response)
{
	char buf[2048] = { 0 };

	for (int i = 0; i < request->values_size(); i++)
	{
		auto value = request->values(i);

		int len = enclave_provider_->encrypt(value.c_str(), (int)value.length(), buf, (int)sizeof(buf));

		if (len == 0)
		{
			return ::grpc::Status(::grpc::StatusCode::CANCELLED, "Encrypt error");
		}

		auto encrypt_value = response->add_values();
		*encrypt_value = std::string(buf, len).c_str();
	}

	return ::grpc::Status::OK;
}

::grpc::Status Service::Decrypt(::grpc::ServerContext* context, const ::enclave::Fields* request, ::enclave::Fields* response)
{
	char buf[2048] = { 0 };

	for (int i = 0; i < request->values_size(); i++)
	{
		auto value = request->values(i);

		int len = enclave_provider_->decrypt(value.c_str(), (int)value.length(), buf, (int)sizeof(buf));

		if (len == 0)
		{
			return ::grpc::Status(::grpc::StatusCode::CANCELLED, "Encrypt error");
		}

		auto decrypt_value = response->add_values();
		*decrypt_value = std::string(buf, len).c_str();
	}

	return ::grpc::Status::OK;
}

std::shared_ptr<EnclaveProvider> Service::get_enclave_porvider(const std::string & id)
{
	CriticalSectionUtil cs_util(cs_);

	if (enclave_providers_map_.find(id) != enclave_providers_map_.end())
	{
		return enclave_providers_map_[id];
	}

	std::shared_ptr<EnclaveProvider> enclave_provider = std::shared_ptr<EnclaveProvider>(new EnclaveProvider(enclave_provider_->enclave_id_));
	enclave_provider->initVault(enclave_provider_->enclave_id_, id);

	enclave_providers_map_[id] = enclave_provider;

	return enclave_provider;
}

